#include "psd2.h"
#include "exectuor.h"

__device__ DuplexSpongeMode DuplexSpongeMode::Absorbing(int next_absorb_index)
{
    DuplexSpongeMode mode;
    mode.type = DuplexSpongeModeType::Absorbing;
    mode.next_absorb_index = next_absorb_index;
    return mode;
}

__device__ DuplexSpongeMode DuplexSpongeMode::Squeezing(int index)
{
    DuplexSpongeMode mode;
    mode.type = DuplexSpongeModeType::Squeezing;
    mode.next_squeeze_index = index;
    return mode;
}

__device__ DuplexSpongeMode::DuplexSpongeMode()
{
    type = DuplexSpongeModeType::Absorbing;
    next_squeeze_index = 0;
}

//Psd2
__device__ Field *psd2_domain = nullptr;
__device__ CuVector<CuVector<Field>> *psd2_mds = nullptr;
__device__ CuVector<CuVector<Field>> *psd2_ark = nullptr;

__device__ Psd2::Psd2()
{
    RATE = 2;
    CAPACITY = 1;
    partial_rounds = 31;
    full_rounds = 8;
    alpha = Field::constant(17);
    m_exectuor = nullptr;
}

__device__ void Psd2::setExectuor(Exectuor* e)
{
    m_exectuor = e;
}

__device__ CuVector<Field> Psd2::hash_many(const CuVector<Field>& input, uint16_t num_outputs) {
    // Construct the preimage: [ DOMAIN || LENGTH(INPUT) || [0; RATE-2] || INPUT ].
    CuVector<Field> preimage;
    preimage.push_back(*psd2_domain);
    preimage.push_back(Field::constant(input.size()));
    preimage.resize(RATE, Field::zero()); // Pad up to RATE.
    preimage.append(input);

    // Initialize a new sponge.
    CuVector<Field> state(RATE + CAPACITY, Field::zero());
    DuplexSpongeMode mode = DuplexSpongeMode::Absorbing(0);

    // Absorb the input and squeeze the output.
    this->absorb(state, mode, preimage);
    return this->squeeze(state, mode, num_outputs);
}

__device__ void Psd2::absorb(CuVector<Field>& state, DuplexSpongeMode &mode, const CuVector<Field>& input) 
{
    bool should_permute;
    int absorb_index;
    if (!input.isEmpty()) 
    {
        // Determine the absorb index.
        if(mode.type == Absorbing)
        {
            if(mode.next_absorb_index == RATE)
            {
                absorb_index = 0;
                should_permute = true;   
            }
            else
            {
                absorb_index = mode.next_absorb_index;
                should_permute = false;
            }
        }
        else
        {
            absorb_index = 0;
            should_permute = false;
        }

        // Proceed to permute the state, if necessary.
        if( should_permute ){
            this->permute(state);
        }

        auto remaining = input;
        while(true){
            // Compute the starting index.
            auto start = CAPACITY + absorb_index;

            // Check if we can exit the loop.
            if( absorb_index + remaining.size() <= RATE) {
                // Absorb the state elements into the input.
                for (int i = 0; i < remaining.size(); i++)
                {
                    auto element = remaining[i];
                    state[start + i] = state[start + i] + element;
                }
                // Update the sponge mode.
                mode = DuplexSpongeMode::Absorbing(absorb_index + remaining.size());
                return;
            }

            // Otherwise, proceed to absorb `(rate - absorb_index)` elements.
            auto num_absorbed = RATE - absorb_index;
            for (int i = 0; i < num_absorbed; i++)
            {
                auto element = remaining[i];
                state[start + i] = state[start + i] + element;
            }

            // Permute the state.
            this->permute(state);

            // Repeat with the updated input slice and absorb index.
            remaining = remaining.mid(num_absorbed);
            absorb_index = 0;
        }
    }
}

/// Squeeze the specified number of state elements into the output.
__device__ CuVector<Field> Psd2::squeeze(CuVector<Field>& state, DuplexSpongeMode &mode, uint16_t num_outputs) {
    CuVector<Field> output(num_outputs,Field::zero());
    if(num_outputs != 0){
        this->squeeze_internal(state, mode, output);
    }
    return output;
}

/// Squeeze the state elements into the output.
__device__ void Psd2::squeeze_internal(CuVector<Field>& state, DuplexSpongeMode &mode, CuVector<Field>& output) {
    int squeeze_index;
    bool should_permute;
    // Determine the squeeze index.
    if (mode.type == Absorbing)
    {
        squeeze_index = 0;
        should_permute = true;
    }
    else 
    {
        if (mode.next_squeeze_index == RATE)
        {
            squeeze_index = 0;
            should_permute = true;
        }
        else
        { 
            squeeze_index = mode.next_squeeze_index;
            should_permute = false;
        }
    };

    // Proceed to permute the state, if necessary.
    if (should_permute) {
        this->permute(state);
    }

    auto remaining = output;
    while(true){
        // Compute the starting index.
        auto start = CAPACITY + squeeze_index;

        // Check if we can exit the loop.
        if (squeeze_index + remaining.size() <= RATE ){
            // Store the state elements into the output.
            remaining = state.mid(start,remaining.size());
            // Update the sponge mode.
            mode = DuplexSpongeMode::Squeezing(squeeze_index + remaining.size());
            output = remaining;
            return;
        }

        // Otherwise, proceed to squeeze `(rate - squeeze_index)` elements.
        auto num_squeezed = RATE - squeeze_index;
        for(int i = 0; i < num_squeezed; i++)
            remaining[i] = state[i + start];

        // Permute.
        this->permute(state);

        // Repeat with the updated output slice and squeeze index.
        remaining = remaining.mid(num_squeezed);
        squeeze_index = 0;
    }
}

/// Apply the additive round keys in-place.
__device__ void Psd2::apply_ark(CuVector<Field>& state, int round) 
{
    for (int i = 0; i < state.size(); i++){
        auto element = state[i];
        state[i] = state[i] + (*psd2_ark)[round][i];
    }
}

/// Apply the S-Box based on whether it is a full round or partial round.
__device__ void Psd2::apply_s_box(CuVector<Field>& state, bool is_full_round) 
{
    if (is_full_round) {
        // Full rounds apply the S Box (x^alpha) to every element of state
        for (int i = 0; i < state.size(); i++) {
            state[i] = m_exectuor->field_pow(state[i],this->alpha);
        }
    }
    else {
        // Partial rounds apply the S Box (x^alpha) to just the first element of state
        state[0] = m_exectuor->field_pow(state[0],this->alpha);
    }
}

/// Apply the Maximally Distance Separating (MDS) matrix in-place.
__device__ void Psd2::apply_mds(CuVector<Field>& state) 
{
    CuVector<Field> new_state;
    for(int i = 0; i < state.size(); i++) 
    {
        auto accumulator = Field::zero();
        for (int j = 0; j < state.size(); j++)
        {
            accumulator = accumulator + state[j] * (*psd2_mds)[i][j];
        }
        new_state.push_back(accumulator);
    }
    state = new_state;
}

/// Apply the permutation for all rounds in-place.
__device__ void Psd2::permute(CuVector<Field>& state) 
{
    // Determine the partial rounds range bound.
    auto full_rounds_over_2 = this->full_rounds / 2;
    auto partial_round_range_s = full_rounds_over_2;
    auto partial_round_range_e = (full_rounds_over_2 + this->partial_rounds);

    //qDebug() << "permute in:" << state;

    // Iterate through all rounds to permute.
    for(int i = 0; i < (this->partial_rounds + this->full_rounds); i++) {
        auto is_full_round = !(i >= partial_round_range_s && i < partial_round_range_e);
        this->apply_ark(state, i);
        //qDebug() << "apply_ark:" << state;

        this->apply_s_box(state, is_full_round);
        //qDebug() << "apply_s_box:" << state;

        this->apply_mds(state);
        //qDebug() << "apply_mds:" << state;
    }
}